import os
from sklearn.model_selection import train_test_split


'''数据划分'''
# 预处理输出地址
data_path = '/root/autodl-tmp/BrainTumorSegmentation1/data_set/BraTS2021/dataset'  # 预处理好的数据的存储地址
# 获取指定路径 data_path 下的所有文件和文件夹的名称，并将它们存储在列表 train_and_test_ids 中
train_and_test_ids = os.listdir(data_path)

train_ids, val_test_ids = train_test_split(train_and_test_ids, test_size=0.2,random_state=21) # 将训练和测试的样本ID划分为训练集和验证集 验证集占总样本的20%
test_ids, val_ids = train_test_split(val_test_ids, test_size=0.5,random_state=21)
# print("Using {} images for training, {} images for validation, {} images for testing.".format(len(train_ids),len(val_ids),len(test_ids)))
# print("Using %d images for training, %d images for validation, %d images for testing." % (len(train_ids), len(val_ids), len(test_ids)))
print(f"Using {len(train_ids)} images for training, {len(val_ids)} images for validation, {len(test_ids)} images for testing.")


# 排序
train_ids.sort()
val_ids.sort()
test_ids.sort()

# 将经过排序的训练集样本ID写入到名为 'train.txt' 的文件中
with open('/root/autodl-tmp/BrainTumorSegmentation1/data_set/BraTS2021/train.txt','w') as f:
    f.write('\n'.join(train_ids))

with open('/root/autodl-tmp/BrainTumorSegmentation1/data_set/BraTS2021/valid.txt','w') as f:
    f.write('\n'.join(val_ids))

with open('/root/autodl-tmp/BrainTumorSegmentation1/data_set/BraTS2021/test.txt','w') as f:
    f.write('\n'.join(test_ids))
